{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resynchronizing Jittery AES Power Traces\n",
    "\n",
    "Supported setups:\n",
    "\n",
    "SCOPES:\n",
    "\n",
    "* OPENADC\n",
    "* CWNANO\n",
    "\n",
    "PLATFORMS:\n",
    "\n",
    "* CWLITEARM\n",
    "* CWLITEXMEGA\n",
    "* CWNANO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens if things aren't as clean as we made them out to be? We can use preprocessing modules!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'TINYAES128C'\n",
    "num_traces = 50\n",
    "CHECK_CORR = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capturing Jittery Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rebuilding New Firmware"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In file `chipwhisperer/hardware/victims/firmware/simpleserial-aes/simpleserial-aes.c` find this:\n",
    "\n",
    "```c\n",
    "uint8_t get_pt(uint8_t* pt)\n",
    "{\n",
    "\ttrigger_high();\n",
    "\taes_indep_enc(pt); /* encrypting the data block */\n",
    "\ttrigger_low();\n",
    "\tsimpleserial_put('r', 16, pt);\n",
    "\treturn 0x00;\n",
    "}\n",
    "```\n",
    "\n",
    "and add some random delay:\n",
    "\n",
    "```c\n",
    "uint8_t get_pt(uint8_t* pt)\n",
    "{\n",
    "\ttrigger_high();\n",
    "       for(volatile uint8_t k = 0; k < (*pt & 0x0F); k++);\n",
    "\taes_indep_enc(pt); /* encrypting the data block */\n",
    "\ttrigger_low();\n",
    "\tsimpleserial_put('r', 16, pt);\n",
    "\treturn 0x00;\n",
    "}\n",
    "```\n",
    "\n",
    "This deterministic delay is NOT a good countermeasure, but is much easier to write in a single line since we don't have a CSPRNG linked in. We'll break the jitter without relying on the deterministic aspect though, so our attack would work against a better jitter source.\n",
    "\n",
    "**Be sure to remove this function afterwards so you don't break your code!**\n",
    "\n",
    "We can build the code (change the platform as needed), and confirm the output of the following works as you expect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../hardware/victims/firmware/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2 EXTRA_OPTS=ADD_JITTER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's go ahead. We'll have to program the file we built, so be sure to confirm we are using the right file!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"Helper_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time\n",
    "\n",
    "fw_path = '../hardware/victims/firmware/simpleserial-aes/simpleserial-aes-{}.hex'.format(PLATFORM)\n",
    "\n",
    "modtime = os.path.getmtime(fw_path)\n",
    "print(\"File build time: {:s} (built {:.2f} mins ago)\".format(str(time.ctime(modtime)), (time.time() - modtime)/60.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, before we capture our traces, we'll need to create a ChipWhipserer project, since that's what Analyzer expects for an input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project = cw.create_project(\"projects/jittertime\", overwrite = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Capturing Traces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below you can see the capture loop. The main body of the loop loads some new plaintext, arms the scope, sends the key and plaintext, then finally records and our new trace into our trace class. We'll also keep track of our keys manually for checking our answer later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Capture Traces\n",
    "from tqdm import tnrange\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "ktp = cw.ktp.Basic()\n",
    "\n",
    "\n",
    "for i in tnrange(num_traces, desc='Capturing traces'):\n",
    "    key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save project file\n",
    "project.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're now done with the ChipWhisperer hardware, so we should disconnect from the scope and target:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup the connection to the target and scope\n",
    "scope.dis()\n",
    "target.dis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fix the jitter, we'll need to add our traces to a preprocessing module. We can feed our project right into `cwa.cpa()`, but we could also add pre-processing inbetween (more about this later)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import chipwhisperer.analyzer as cwa\n",
    "project = cw.open_project(\"projects/jittertime\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time we're going to do a few things. First we will get the traces, and plot a few of them as-is. You can adjust the traces plotted by adjusting the `range(10)`. For example `range(1)` plots the first trace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "from bokeh.palettes import Dark2_5 as palette\n",
    "import itertools  \n",
    "\n",
    "output_notebook()\n",
    "p = figure(sizing_mode='scale_width', plot_height=300)\n",
    "\n",
    "# create a color iterator\n",
    "colors = itertools.cycle(palette)  \n",
    "\n",
    "x_range = range(0, len(project.waves[0]))\n",
    "for i, color in zip(range(5), colors): #Adjust range(n) to plot certain traces\n",
    "    p.line(x_range, project.waves[i], color=color)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So how do we fix that? To begin with, you should plot only a single trace to make your life more clear. You'll need to figure out a very unique area. For example see the following figure showing a single plot. In this example the location of         \n",
    "**A** is unique, but B would have many matches within that same trace, even nearby:\n",
    "![Resync example trace](img/resync_goodbad.png)\n",
    "\n",
    "We will specify two items:\n",
    "- A window with the \"unique\" area defined.\n",
    "- How far we will shift the window (+/- points) to search for the best match.\n",
    "\n",
    "You can use the following code to define the `target_window` and `max_shift`. Try a few values until you find something that works. Also try some poor example, and also try plotting more traces to confirm your match is working in real life."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resync_traces = cwa.preprocessing.ResyncSAD(project)\n",
    "resync_traces.ref_trace = 0\n",
    "\n",
    "if PLATFORM == \"CWNANO\":\n",
    "    #Define a target window here. 500,900 for example is good based on above. But try some different values.\n",
    "    resync_traces.target_window = (300, 700)\n",
    "\n",
    "    # Define max_shift. Must not cause target_window to go outside of valid data. Try 16-600 range. Ideal value varies with how\n",
    "    # much jitter is in original data. \n",
    "    resync_traces.max_shift = 300\n",
    "elif PLATFORM == \"CWLITEXMEGA\" or PLATFORM == \"CW303\":\n",
    "    #Define a target window here. 500,900 for example is good based on above. But try some different values.\n",
    "    resync_traces.target_window = (1000, 1400)\n",
    "\n",
    "    # Define max_shift. Must not cause target_window to go outside of valid data. Try 16-600 range. Ideal value varies with how\n",
    "    # much jitter is in original data. \n",
    "    resync_traces.max_shift = 1000\n",
    "elif PLATFORM==\"CWLITEARM\":\n",
    "    #Define a target window here. 500,900 for example is good based on above. But try some different values.\n",
    "    resync_traces.target_window = (700, 1500)\n",
    "\n",
    "    # Define max_shift. Must not cause target_window to go outside of valid data. Try 16-600 range. Ideal value varies with how\n",
    "    # much jitter is in original data. \n",
    "    resync_traces.max_shift = 700\n",
    "\n",
    "#Uses objects from previous cells (plotting etc), so \n",
    "output_notebook()\n",
    "p = figure()\n",
    "new_proj = resync_traces.preprocess()\n",
    "\n",
    "for i, color in zip(range(0, 5), colors):\n",
    "    p.line(x_range, new_proj.waves[i], color=color)\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If this all works - let's just continue the attack! Do so as below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(new_proj, leak_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then actually run it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cb = cwa.get_jupyter_callback(attack)\n",
    "attack_results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You *should* see the PGE reach 0 for each byte. If not, you might need to adjust the SAD resync. You could also need to increase the length of the sample capture for example. You may notice that it starts working OK and then fails, due to later traces become unsychronized."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting Correlation Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "\n",
    "attack_results = attack.results\n",
    "plot_data = cwa.analyzer_plots(attack_results)\n",
    "bnum = 0\n",
    "\n",
    "ret = plot_data.output_vs_time(bnum)\n",
    "\n",
    "output_notebook()\n",
    "p = figure()\n",
    "p.line(ret[0], ret[2], line_color='green')\n",
    "p.line(ret[0], ret[3], line_color='green')\n",
    "\n",
    "p.line(ret[0], ret[1], line_color='red')\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see a graph of red and green in time (samples). In red is the correlation of the correct subkey for the first byte, while the rest are in green.\n",
    "\n",
    "You should see two or three distinctive red spikes. The first is the spot where the sbox lookup for the subkey we guessed actually happens (the later ones are from later steps in the AES operation).\n",
    "\n",
    "What about the rest of the bytes in the key? We can get and plot that easily as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rets = []\n",
    "for i in range(0, 16):\n",
    "    rets.append(plot_data.output_vs_time(i))\n",
    "\n",
    "p = figure()\n",
    "for ret in rets:\n",
    "    p.line(ret[0], ret[2], line_color='green')\n",
    "    p.line(ret[0], ret[3], line_color='green')\n",
    "    \n",
    "for ret in rets:\n",
    "    p.line(ret[0], ret[1], line_color='red')\n",
    "\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome! You should have now completed a resynchronization of power traces. This is a very useful tool, and you can see how making a simple class could extend this work."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = list(project.keys[0])\n",
    "recv_key = [kguess[0][0] for kguess in attack_results.find_maximums()]\n",
    "assert key == recv_key, \"Failed to recover encryption key\\nGot: {}\\nExpected: {}\".format(recv_key, key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (attack_results.pge == [0]*16), \"PGE for some bytes not zero: {}\".format(attack_results.pge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if CHECK_CORR:\n",
    "    max_corrs = [kguess[0][2] for kguess in attack_results.find_maximums()]\n",
    "    assert (np.all([corr > 0.75 for corr in max_corrs])), \"Low correlation in attack (corr <= 0.75): {}\".format(max_corrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
